from PIL import Image
from torchfusion_utils.fp16 import convertToFP16
from torchfusion_utils.initializers import *
from torchfusion_utils.metrics import Accuracy
from torchfusion_utils.models import load_model, save_model

import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms, models
import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import Variable
from tqdm import tqdm

class CNNtrain:
    def __init__(self):
        self.CLS2ID = {}
        self.ID2CLS = {}
        self.num_classes = 0

    # def image_display(self, image, title,fontsize):
    #     image = image / 2 + 0.5
    #     numpy_image = image.numpy()
    #     transposed_numpy_image = np.transpose(numpy_image, (1, 2, 0))
    #     plt.figure(figsize=(20, 4))
    #     plt.imshow(transposed_numpy_image)
    #     plt.yticks([])
    #     plt.xticks([])

    #     plt.title(title, fontsize=fontsize)
    #     plt.show()

    def load_data(self, train_data_dir, test_data_dir, batch_size=64):
        transforms_train = transforms.Compose([
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        transforms_test = transforms.Compose([
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        train_data = datasets.ImageFolder(root=train_data_dir, transform=transforms_train)
        test_data = datasets.ImageFolder(root=test_data_dir, transform=transforms_test)

        self.CLS2ID = train_data.class_to_idx
        self.ID2CLS = {v: k for k, v in self.CLS2ID.items()}
        self.num_classes = len(self.CLS2ID)

        # 打印类别信息用于调试
        print(f"检测到的类别数量: {self.num_classes}")
        print(f"训练集类别: {train_data.classes}")
        print(f"类别到索引映射: {self.CLS2ID}")
        print(f"索引到类别映射: {self.ID2CLS}")

        train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
        test_data_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=True)
        return train_data_loader, test_data_loader

    def train(self, train_data_loader, test_data_loader, n_epochs, save_path=None):
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        
        # 使用预训练模型，但修改最后一层
        print(f"创建模型，类别数: {self.num_classes}")
        model = models.resnet50(pretrained=True)  # 加载预训练权重
        # 修改全连接层，适应你的类别数量
        model.fc = nn.Linear(model.fc.in_features, self.num_classes)
        model = model.to(device)
        
        criteria = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

        milestones = [100, 150]
        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1)
        train_acc = Accuracy()
        validation_acc = Accuracy(topK=1)
        training_loss_array = []
        validation_loss_array = []

        saving_criteria_of_model = 0

        for epoch in range(n_epochs):
            total_test_data = 0
            training_loss = 0
            validation_loss = 0

            train_acc.reset()

            for data, target in tqdm(train_data_loader, desc="Training"):
                data, target = data.to(device), target.to(device)
                optimizer.zero_grad()
                predictions = model(data)
                loss = criteria(predictions, target)
                loss.backward()
                optimizer.step()

                training_loss += loss.item() * data.size(0)
                train_acc.update(predictions, target)

            scheduler.step()

            with torch.no_grad():
                validation_acc.reset()
                for data, target in test_data_loader:
                    data, target = data.to(device), target.to(device)
                    predictions = model(data)
                    loss = criteria(predictions, target)
                    validation_acc.update(predictions, target)
                    total_test_data += target.size(0)
                    validation_loss += loss.item() * data.size(0)

            training_loss = training_loss / len(train_data_loader.dataset) 
            validation_loss = validation_loss / total_test_data 
            training_loss_array.append(training_loss)
            validation_loss_array.append(validation_loss)
            print(f'{epoch + 1} / {n_epochs} Training loss: {training_loss}, Train Accuracy: {train_acc.getValue()}, Validation loss: {validation_loss}, Validation Accuracy: {validation_acc.getValue()}')
            if save_path != None:
                if saving_criteria_of_model < validation_acc.getValue():
                    # 保存模型时同时保存类别映射信息
                    checkpoint = {
                        'model_state_dict': model.state_dict(),
                        'CLS2ID': self.CLS2ID,
                        'ID2CLS': self.ID2CLS,
                        'num_classes': self.num_classes
                    }
                    torch.save(checkpoint, save_path)
                    saving_criteria_of_model = validation_acc.getValue()
                    print(f"模型已保存到: {save_path}")
        
        plt.figure(figsize=(10, 6))
        x_axis = range(n_epochs)
        plt.plot(x_axis, training_loss_array, 'r', validation_loss_array, 'b')
        plt.title('A graph of training loss vs validation loss')
        plt.legend(['train loss', 'validation loss'])
        plt.xlabel('Number of Epochs')
        plt.ylabel('Loss')
        plt.show()
        return model, saving_criteria_of_model

    def image_display(self, img_path, title, fonsize):
        plt.figure(figsize=(10, 6))
        
        plt.imshow(Image.open(img_path))
        plt.title(title, fontsize=fonsize)
        plt.yticks([])
        plt.xticks([])
        plt.show()

    def load_model_with_mapping(self, model_path):
        """加载模型并恢复类别映射"""
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        
        try:
            # 尝试加载检查点
            checkpoint = torch.load(model_path, map_location=device)
            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                # 从检查点加载
                model = models.resnet50(pretrained=False)
                model.fc = nn.Linear(model.fc.in_features, checkpoint['num_classes'])
                model.load_state_dict(checkpoint['model_state_dict'])
                # 恢复类别映射
                self.CLS2ID = checkpoint['CLS2ID']
                self.ID2CLS = checkpoint['ID2CLS']
                self.num_classes = checkpoint['num_classes']
                model = model.to(device)
                print("从检查点加载模型和类别映射")
                print(f"类别数量: {self.num_classes}")
                print(f"类别映射: {self.ID2CLS}")
                return model
            else:
                # 旧格式，直接加载模型
                print("加载旧格式模型，可能需要手动设置类别映射")
                return checkpoint.to(device)
        except Exception as e:
            print(f"模型加载失败: {e}")
            return None

    def model_inference_results(self, img_path, load_saved_model):
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        transformer = transforms.Compose([
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])
        img = Image.open(img_path)
        img_processed = transformer(img).unsqueeze(0)
        img_var = Variable(img_processed, requires_grad=False)
        img_var = img_var.to(device)
        load_saved_model.eval()
        
        with torch.no_grad():
            logp = load_saved_model(img_var)
            expp = torch.softmax(logp, dim=1)
            
            # 获取完整的置信度分布
            confidence_distribution = expp.squeeze().detach().cpu().numpy()
            
            # 打印完整的置信度分布
            print("=== 完整置信度分布 ===")
            for i, confidence in enumerate(confidence_distribution):
                class_name = self.ID2CLS.get(i, f"未知类别{i}")
                print(f"{class_name}: {confidence:.4f} ({confidence*100:.2f}%)")
            
            # 获取最高置信度的结果
            confidence, clas = expp.topk(1, dim=1)
            confidence = round(confidence.item(), 4)
            clas = clas.item()
            
            # 解释结果
            print("\n=== 结果解释 ===")
            print(f"模型最终判断为: {self.ID2CLS.get(clas, f'未知类别{clas}')}")
            print(f"最高置信度: {confidence:.4f} ({confidence*100:.2f}%)")
            
            # 分析置信度分布特点
            sorted_confidences = sorted(confidence_distribution, reverse=True)
            if len(sorted_confidences) > 1:
                confidence_gap = sorted_confidences[0] - sorted_confidences[1]
                print(f"与第二可能类别的置信度差距: {confidence_gap:.4f}")
                
                if confidence > 0.8:
                    print("模型判断非常确定")
                elif confidence > 0.6:
                    print("模型判断比较确定")
                elif confidence > 0.4:
                    print("模型判断相对确定")
                else:
                    print("模型判断不太确定，可能存在混淆")
            
            # 可视化置信度分布
            plt.figure(figsize=(10, 6))
            classes = [self.ID2CLS.get(i, f"Class{i}") for i in range(len(confidence_distribution))]
            colors = ['lightcoral' if i == np.argmax(confidence_distribution) else 'skyblue' for i in range(len(confidence_distribution))]
            
            plt.bar(classes, confidence_distribution, color=colors, alpha=0.7)
            plt.title('各类别置信度分布', fontsize=15)
            plt.xlabel('类别', fontsize=12)
            plt.ylabel('置信度', fontsize=12)
            plt.xticks(rotation=45)
            plt.grid(True, alpha=0.3)
            
            # 在柱状图上显示数值
            for i, v in enumerate(confidence_distribution):
                plt.text(i, v + 0.01, f'{v:.3f}', ha='center', va='bottom', fontsize=10)
            
            plt.tight_layout()
            plt.show()
            
            return confidence, clas, confidence_distribution

    def visualize_confidence_distribution(self, confidence_distribution, title="置信度分布"):
        """单独的可视化置信度分布方法"""
        plt.figure(figsize=(10, 6))
        classes = [self.ID2CLS.get(i, f"Class{i}") for i in range(len(confidence_distribution))]
        colors = ['lightcoral' if i == np.argmax(confidence_distribution) else 'skyblue' for i in range(len(confidence_distribution))]
        
        plt.bar(classes, confidence_distribution, color=colors, alpha=0.7)
        plt.title(title, fontsize=15)
        plt.xlabel('类别', fontsize=12)
        plt.ylabel('置信度', fontsize=12)
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)
        
        # 在柱状图上显示数值
        for i, v in enumerate(confidence_distribution):
            plt.text(i, v + 0.01, f'{v:.3f}', ha='center', va='bottom', fontsize=10)
        
        plt.tight_layout()
        plt.show()